import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from torch.utils.data import DataLoader, TensorDataset
import matplotlib.pyplot as plt

dataset = load_iris()

x = torch.tensor(dataset['data'], dtype=torch.float32)
y = torch.tensor(dataset['target'], dtype=torch.long)

x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.4)

train_dataset = TensorDataset(x_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)



class IRIS(nn.Module):
    def __init__(self):
        super(IRIS, self).__init__()
        self.fc1 = nn.Linear(4 , 16)
        self.fc2 = nn.Linear(16 , 3)


    def forward(self, x):
        x = torch.tanh(self.fc1(x))
        x = self.fc2(x)
        return x

model = IRIS()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adamax(model.parameters(), lr=0.08)


num_epochs = 60
for epoch in range(num_epochs):
    for i, (inputs, labels) in enumerate(train_loader):
        optimizer.zero_grad()
        output = model(x_train)
        loss = criterion(output, y_train)
        loss.backward()
        optimizer.step()
        print(f'Epoch [{epoch+1}/{num_epochs}], Loss: {loss.item()}')


# 在模型上进行预测
with torch.no_grad():
    model.eval()
    y_pred = model(x_test)
    _, predicted = torch.max(y_pred, 1)

# 计算准确率
correct = (predicted == y_test).sum().item()
total = y_test.size(0)
accuracy = correct / total * 100
print(f'Accuracy on test set: {accuracy:.4f}%')

